"""
落点预测模块

date: 2025/7/15
author: SiHeng Tang
file: drop_predict.py
copyright(c) DFSA Eastern Software
此程序不能提供任何担保 WITHOUT WARRANTY OF ANY KIND
"""
import logging
import time
from copy import deepcopy

from async_sche import TaskWrapper, BaseTaskGen, ExitGenerationError
from helpers.drop_physics import predict_drop_point
from helpers.geo_math_trans import GeoVector
from mengji import Vehicle
from mod_loader import ConfigHandler
from task_element import DataEndpoint, DataPack

logger = logging.getLogger(__name__)
handler = ConfigHandler(__name__)

DEFAULT_CONFIG = {
    "label": "drop_predictor",
    "model": "simple_g",
    "default_enable": 0
}


class DropPredictor(BaseTaskGen):
    """
    落点预测模块使用工具函数库中的落点模型接口进行落点计算。
    此模块的预测数据通过端点发送，可以使用 ``drop_point_ep`` 监听，
    端点荷载类型为 GeoVector。
    """

    def __init__(self, config: dict):
        super().__init__(priority=TaskWrapper.IDLE_PRIORITY + 10)
        self._drop_model = config["model"]
        self._enable_drop_predict = config["default_enable"]

        self.recv_gps = DataEndpoint()
        self.pub_drop_point_predicted_ep = DataEndpoint()

        self._last_act_time = time.time()

    def set_enable_drop_predict(self, enable=True):
        self._enable_drop_predict = enable

    def _on_gen(self):
        if not (self._enable_drop_predict and self.recv_gps.detect_new_in_dp()):
            raise ExitGenerationError()

    async def _func(self):
        pos_dict, src_time = self.recv_gps.get_dp().get_s("payload", "pack_time")

        # 数据合法，相对高度为正，地面上可能波动到负数
        if not (pos_dict and pos_dict["relative_alt"] > 0.):
            return

        origin_delta = predict_drop_point(alt=pos_dict["relative_alt"],
                                          velocity=[pos_dict["vx"], pos_dict["vy"], pos_dict["vz"]],
                                          model=self._drop_model)

        # 注意从 AP “北东地” 换到 GV “东北天”
        src_position = GeoVector(lat=pos_dict["lat"], lon=pos_dict["lon"])
        drop_point = src_position.add_difference(GeoVector(dx=origin_delta[1], dy=origin_delta[0]))

        self.pub_drop_point_predicted_ep.set_dp(DataPack(payload=drop_point,
                                                         src_time=src_time))

    def _on_exit(self):
        self.recv_gps = None
        self.pub_drop_point_predicted_ep = None


@handler.add_handler("unit_proc:drop_predict")
def load_config(veh: Vehicle, cfg_section: dict):
    logger.debug(f"drop predict load config {cfg_section}")

    cfg = deepcopy(DEFAULT_CONFIG)
    cfg.update(cfg_section)

    veh.reg_part(cfg["label"], DropPredictor(cfg), is_sche=True)
